#ifndef __MULTI_DISPATCHER_H__
#define __MULTI_DISPATCHER_H__

#include <functional>
#include <unordered_map>
#include <utility>
#include <vector>
#include "collab/utmutex.h"
namespace utools::collab
{
    template <typename KeyT, typename FuncT>
    class MultiDispatcher;

    template <typename KeyT, typename ReturnType, typename... Args>
    class MultiDispatcher<KeyT, ReturnType(Args...)>
    {
    public:
        using WorkFunctionType = std::function<ReturnType(Args...)>;

        MultiDispatcher() = default;
        virtual ~MultiDispatcher() = default;

        template <typename Func>
        MultiDispatcher &map(const KeyT key, Func &&func)
        {
            collab::UniqueLock lock(__mutex);
            __dispatcher_map[key].emplace_back(std::forward<Func>(func));
            return *this;
        }

        template <typename T, typename MemberFunc>
        MultiDispatcher &map(const KeyT key, MemberFunc &&func, T *obj)
        {
            collab::UniqueLock lock(__mutex);
            __dispatcher_map[key].emplace_back(
                [obj, func](Args... args)
                { return (obj->*func)(std::forward<Args>(args)...); });
            return *this;
        }

        /// @brief 删除映射调度任务
        /// @param key 绑定的键
        /// @return 返回任务对象本身
        MultiDispatcher &unmap(const KeyT key)
        {
            collab::UniqueLock lock(__mutex);
            __dispatcher_map.erase(key);
            return *this;
        }

        /// @brief 调度任务
        /// @param key 调度任务标识
        /// @param args 可变参数
        ReturnType dispatch(const KeyT key, Args... args)
        {
            collab::UniqueLock lock(__mutex);
            auto it = __dispatcher_map.find(key);
            if (it != __dispatcher_map.end())
            {
                for (auto &func : it->second)
                {
                    func(std::forward<Args>(args)...);
                }
            }
            return ReturnType();
        }

    private:
        std::unordered_map<KeyT, std::vector<WorkFunctionType>> __dispatcher_map;
        collab::Mutex __mutex;
    };
} // namespace utools::utcollab

#endif
